#ifndef EST_POLYA_NEWTON_HEADER
#define EST_POLYA_NEWTON_HEADER

#include<string>
#include<fstream>
#include<sstream>
#include<iostream>
#include<algorithm>
#include<iterator>
#include<vector>
#include<map>
#include<cstdlib>
#include<ctime>
#include<cmath>
#include <gsl/gsl_sort.h>
#include <gsl/gsl_statistics.h>
#include <gsl/gsl_sf_psi.h>

#include "def.hpp"
#include "utils.hpp"
#include "io.hpp"
#include "dist_polya.hpp"

using namespace std;

void est_polya_newton(int**& data,int*& n_i, double*& param, int& K, int& N){
  double* g = new double[K];
  double* q = new double[K];
  double* fake_param = new double[K];
  double* step = new double[K];
  double z = 0.0;
  double b = 0.0;
  double temp = 0.0;
  double temp_1 = 0.0;
  double temp_2 = 0.0;

  //  double likelihood = 0.0;
  double fake_likelihood = 0.0;

  int ITER = 0;

  double reg = 0.1; // the regularization step

  while(1){
    double ParamSum = get_sum(param, K);
    double likelihood = polya_likelihood(data,n_i,param, K,N);
    //    cout << "ITER :" << ITER << " Likelihood :" << likelihood << " --- REG:" << reg << endl;
    // clean
    while(1){
      //      cout << reg << endl;
      temp = 0.0;
      temp_1 = 0.0;
      temp_2 = 0.0;
      z = 0.0;
      for(int k =0; k < K; k++){
	g[k] = 0.0;
	q[k] = 0.0;
      }

      for(int i =0; i < N; i++){
	temp += gsl_sf_psi(n_i[i] + ParamSum);
	for(int k=0; k < K; k++){
	  g[k] = g[k] + gsl_sf_psi(data[i][k]+param[k]);
	  q[k] = q[k] + gsl_sf_psi_n(1,data[i][k] + param[k]) - gsl_sf_psi_n(1,param[k]);
	}
	z = z + gsl_sf_psi_n(1,ParamSum) - gsl_sf_psi_n(1,n_i[i] + ParamSum);      
      }

      for(int k = 0; k < K; k++){
	g[k] = g[k] + N * (gsl_sf_psi(ParamSum) - gsl_sf_psi(param[k])) - temp;
	g[k] = g[k] / N;
	q[k] = q[k] / N;
	q[k] = q[k] - reg; // regularization
	temp_1 = temp_1 + g[k] / q[k];
	temp_2 = temp_2 + 1.0 / q[k];
      }
      z = z / N;
      b = temp_1 / (1.0 / z + temp_2);

      // prepare update
      int update_flag = 0;
      for(int k=0; k < K; k++){
	step[k] = (g[k] - b)/q[k];
	if(param[k] - step[k] > 0)
	  update_flag += 1;
      }
      if(update_flag == K){
	for(int k =0; k< K; k++){
	  fake_param[k] = param[k] - step[k];
	}
	fake_likelihood = polya_likelihood(data,n_i,fake_param, K,N);
	if(fake_likelihood > likelihood){
	  for(int k =0; k< K; k++){
	    param[k] = fake_param[k];
	    //    cout << param[k] << " ";
	  }
	  //	  cout << endl;
	  reg = reg / 10;
	  break;
	}
      }
      reg = reg * 10;
      if(reg > 1e+6)
	break;
    }

    if(fake_likelihood - likelihood < 1e-5)
      break;
    ITER += 1;
    if(ITER >= 100)
      break;
  }
}

#endif
